# -*- coding: latin-1 -*-
"""
FFT maximum spot mapping visualization in DigitalMicrograph (DM) using the BenMillerScripts package.

Purpose:
- Compute local FFTs over a scanning grid on the front (active) DM image, identify the brightest FFT spot at each
  position, and visualize its angle ("theta") as a color-mapped RGB image.
- Optionally display the intermediate 4D-like data cube used for the analysis (post-binning, FFT, and masking).
- Draw a color scale that indicates how hue maps to angle.

Inputs/Assumptions:
- The front DM image (DM.GetFrontImage()) is the dataset to analyze (A 2D high resolution image, with lattice fringes).
- The script relies on benmillerscripts.FFTArrayAnalysis (FAA) for analysis and DM image creation/display.
- The user may adjust parameters in the "User-adjustable parameters" section below.
- Dimension calibration uses the x-dimension calibration of the input image to set the output map calibration spacing.

Outputs:
- An RGB DM image showing orientation ("theta") of the dominant FFT spot from each FFT computed from overlapping windows. 
  The output image is shown in DM and named for clarity.
- Optional display of the intermediate data cube used in the analysis (controlled by 'show_cube').
- A color scale indicating the mapping between color and maximum spot angle/intensity. 
  (Color Hue corresponds to angle, Color Brightness to intensity)

Notes on calibrations and tags:
- Sets the output map’s dimension calibration based on the input image’s pixel scale, binning and FAA.spacing.
- Copies the source image’s tags into a "Copied Tags" group on the final map.

Libraries:
- Uses benmillerscripts.FFTArrayAnalysis for analysis and DM image creation.

Script Written by Ben Miller
Last Updated 2026-09-01
"""

import benmillerscripts.FFTArrayAnalysis as FAA
import importlib
# Reload analysis module to ensure the latest code is used
FAA = importlib.reload(FAA)

# --------------------------------------------------------------------------
# User-adjustable parameters (kept identical in behavior to original code).
# --------------------------------------------------------------------------

# Variable to encode in color: 'theta' (angle of brightest FFT spot) or 'radius' (spacing of brightest FFT spot)
    #radius produces a valid map, but the color scale is not generated properly, and throws an error
map_var = 'theta'  # Default: 'theta'

# Percentage of the central mask to apply to FFTs (0-100). Higher values mask more of the low-frequency center.
FAA.maskP = 20  # Default: 20

# Width (in pixels) of an additional center mask around the horizontal and vertical center-lines (implementation detail inside FAA).
FAA.maskC_width = 1

# Pre-analysis binning factor (integer > 0). Higher values speed processing but can remove high-frequency detail.
FAA.binning = 2  # Default: 2

# Spacing of the FFT sampling grid across the image (smaller => denser map, longer runtime).
FAA.spacing = 32  # Default: 32

# Size of individual FFT windows (pixels). Smaller => faster but coarser angular resolution.
FAA.FFTsize = 128  # Default: 128


# If True, show the intermediate data cube (post-FFT and masking) akin to a 4D-STEM dataset.
show_cube = True

# If True, pad the map border to mimic edge distances; if False, do not pad.
# An unpadded map has the correct number of pixels to drag the picker tool from the 4D datacube to the map
# A padded map visually matches the original image field of view
FAA.pad_border = False # Default: False

# --------------------------------------------------------------------------
# Acquire input image and run analysis.
# --------------------------------------------------------------------------
# The front DM image is analyzed. FAA.processframe returns:
#   processedframe      -> NumPy-like array or structure used to make the RGB map (handled by FAA.CreateDM_RGB)
#   direction_image     -> DM image of local FFT peak angles (not displayed here)
#   intensity_image     -> DM image of local FFT peak intensities (not displayed here)
#   spacing_image       -> DM image of local FFT peak spacings (not displayed here)
#   diffractogram_max   -> DM image of the global diffractogram pixel-by-pixel maximum (not displayed here)
image = DM.GetFrontImage()
image_name = image.GetName()

(processedframe,
 direction_image,
 intensity_image,
 spacing_image,
 diffractogram_max) = FAA.processframe(
    image,
    map_var,
    show_cube,
    scale=1,
    overlay=False,
    im_data=None,
    RGB_scale_max=None
)

# --------------------------------------------------------------------------
# Create and display the RGB visualization of the selected variable.
# --------------------------------------------------------------------------
# Retrieve the x-axis calibration from the input image. This is used below to set the output map calibration.
# Note: DM's GetDimensionCalibration(axis, 0) returns (origin, scale, unit). Here we use axis=1 for x-scale.
origin, x_scale, scale_unit = image.GetDimensionCalibration(1, 0)

# Convert the processed frame to a DM RGB image using FAA's helper function.
DM_RGB = FAA.CreateDM_RGB(processedframe)

# Set output map dimension calibration. The map grid spacing is FAA.spacing pixels apart relative to the input x-scale.
# We apply the same calibration to both x (axis 1) and y (axis 0) of the output map for a square sampling grid.
DM_RGB.SetDimensionCalibration(1, 0, x_scale * FAA.spacing*FAA.binning, scale_unit, 0)
DM_RGB.SetDimensionCalibration(0, 0, x_scale * FAA.spacing*FAA.binning, scale_unit, 0)

# Name and display the result.
DM_RGB.SetName("Maximum Spot Map of " + image_name)
DM_RGB.ShowImage()

# --------------------------------------------------------------------------
# Copy source tags into a "Copied Tags" subgroup on the current front image.
# --------------------------------------------------------------------------
tg = DM.GetFrontImage().GetTagGroup()
tg_source = image.GetTagGroup()
tg.SetTagAsTagGroup("Copied Tags", tg_source.Clone())

# --------------------------------------------------------------------------
# Optional SI metadata tagging.
# --------------------------------------------------------------------------
DM.DoEvents()
if (FAA.pad_border == False):
    scale = 1
    # Create "SI:Acquisition" tag group if not present.
    if tg.GetTagAsTagGroup("SI:Acquisition") == None:
        tg.SetTagAsTagGroup("SI:Acquisition", DM.NewTagGroup())
    # Set the scale tag value.
    tg.SetTagAsLong("SI:Scale", scale)

# Example of creating an image from other outputs of FAA.processframe.
# Important: If you create a DM image from a NumPy array, always pass array.copy() to avoid view-backed images.
# DM.CreateImage(spacing_image.copy()).ShowImage()

# Clean up image references.
del DM_RGB
del image

# --------------------------------------------------------------------------
# Draw the color scale for the mapping variable.
# --------------------------------------------------------------------------
# This provides a legend indicating the hue mapping for orientation ("theta") or spacing ("radius"). 
# Radius produces a valid maximum spot map, but this color scale is not generated properly, and throws an error
FAA.DrawColorScale(map_var=map_var)